# coding=utf-8
from prompts.problem_classification_prompt import prompt_problem_classification, prompt_content_classification
from utils.functions import Get_List, innovation_dict, motherboard_dict, entrepreneurship_dict
from utils.generate_token import api_key
from langchain_core.output_parsers import StrOutputParser
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from openai import OpenAI, AsyncOpenAI
from collections import OrderedDict
from pprint import pprint
import httpx
import json
import time
import ast

# 招股书切片的问题分类
class Bubble:

    def __init__(self):
        self.llm = ChatZhipuAI(
            api_key=api_key,
            model_name="glm-4-plus",
            streaming=True,
            temperature=0.01,
            top_p=0.01,
        )

        self.client = AsyncOpenAI(
            api_key="5db64dc0-73d7-4139-b755-5bd5483327b4",  # lz
            base_url="https://ark.cn-beijing.volces.com/api/v3",
        )


    async def title_bubble_zp(self, content_text: str, sector: str):
        '''
        :param content_text: 招股书切片
        :param sector: 板块
        :return:
        '''

        s_time = time.time()
        if sector == '科创板':
            problem_type_dict = innovation_dict
        elif sector == '沪市主板' or sector == '深市主板':
            problem_type_dict = motherboard_dict
        elif sector == '创业板':
            problem_type_dict = entrepreneurship_dict

        key_value_list = list(problem_type_dict.items())  # 将字典转换为键值对的列表

        query_message_list = []
        for i in range(0, len(key_value_list), 10):
            batch = key_value_list[i: i + 10]
            type_dict = dict(batch)
            query_message = prompt_content_classification(content_text, type_dict)  # 招股书片段

            # print(1111111111111111111111111111111111111111)
            # print(query_message)

            query_message_list.append(query_message)

        chat_prompt = ChatPromptTemplate.from_messages(
            [("system", "你是一个证券交易所中审核上市文件的专家。"),
             ("human", "{query_message}")]
        )
        chain = chat_prompt | self.llm | StrOutputParser()

        try:
            problem_type_index_list = []
            for chain_results in await chain.abatch(
                    [{"query_message": query_message} for query_message in query_message_list]):
                print(chain_results)
                type_list = eval(Get_List(chain_results))
                problem_type_index_list += type_list
            problem_type_index_list = sorted(list(set(problem_type_index_list)))
            print('problem_type_index_list:', problem_type_index_list)

            if len(problem_type_index_list) == 0:
                print('problem_type_list:', problem_type_index_list)
                return ['该招股书切片不适合被询问']
            else:
                if max(problem_type_index_list) <= 58 and min(problem_type_index_list) >= 0:
                    print('problem_type_list:', problem_type_index_list)
                else:
                    problem_type_index_list = []
                    raise ValueError('问题分类索引不在[0,58]范围内')
        except Exception as e:
            problem_type_index_list = []
            print('Exception:', e)

        problem_type_name_list = []
        for type_index in problem_type_index_list:
            problem_type_name_list.append(problem_type_dict[type_index])
        print(problem_type_name_list)
        # problem_type_name_list = str(problem_type_name_list)

        e_time = time.time()
        print('单条时间', e_time - s_time)

        return problem_type_name_list

    async def title_bubble_v3(self, content_text: str, sector: str):
        '''
        :param content_text: 招股书切片
        :param sector: 板块
        :return:
        '''

        s_time = time.time()
        if sector == '科创板':
            problem_type_dict = innovation_dict
        elif sector == '沪市主板' or sector == '深市主板':
            problem_type_dict = motherboard_dict
        elif sector == '创业板':
            problem_type_dict = entrepreneurship_dict

        query_message = prompt_content_classification(content_text, problem_type_dict)
        headers = {
            'Authorization': 'Bearer 5db64dc0-73d7-4139-b755-5bd5483327b4',  # lz
        }

        json_data = {
            'messages': [
                {
                    'content': f'{query_message}',
                    'role': 'system',
                },
            ],
            'model': 'deepseek-v3-241226',
            "temperature": 0,
            "top_p": 0
        }

        url = 'https://ark.cn-beijing.volces.com/api/v3/chat/completions'  # 确保 URL 是正确的

        while True:
            try:
                async with httpx.AsyncClient() as client:
                    response = await client.post(url, headers=headers, json=json_data, timeout=30)
                    response.raise_for_status()  # 检查响应状态码是否为 200
                    data = response.json()
                    content = data['choices'][0]['message']['content']
                    print(content)
                    content = ast.literal_eval(content)
                    if isinstance(content, list):  # 检查是否为列表
                        break
            except Exception as e:
                print('连接超时或输出非列表，再来一次')
                pass
        try:
            problem_type_index_list = sorted(list(set(content)))
            print('problem_type_index_list:', problem_type_index_list)

            if len(problem_type_index_list) == 0:
                print('problem_type_list:', problem_type_index_list)
                problem_type_index_list = []
                print('该招股书切片不适合被询问')
                # return ['该招股书切片不适合被询问']
            else:
                if max(problem_type_index_list) <= 58 and min(problem_type_index_list) >= 0:
                    print('problem_type_list:', problem_type_index_list)
                else:
                    problem_type_index_list = []
                    raise ValueError('问题分类索引不在[0,58]范围内')
        except Exception as e:
            problem_type_index_list = []
            print('Exception:', e)

        problem_type_name_list = []
        for type_index in problem_type_index_list:
            problem_type_name_list.append(problem_type_dict[type_index])
        print(problem_type_name_list)
        # problem_type_name_list = str(problem_type_name_list)

        e_time = time.time()
        print('单条时间', e_time - s_time)

        return problem_type_name_list

    async def title_bubble_r1(self, content_text: str, sector: str):
        '''
        :param content_text: 招股书切片
        :param sector: 板块
        :return:
        '''

        s_time = time.time()
        if sector == '科创板':
            problem_type_dict = innovation_dict
        elif sector == '沪市主板' or sector == '深市主板':
            problem_type_dict = motherboard_dict
        elif sector == '创业板':
            problem_type_dict = entrepreneurship_dict

        try:
            print("----- streaming request -----")
            query_message = prompt_content_classification(content_text, problem_type_dict)  # 招股书片段
            print('query_message')
            print(query_message)

            stream = await self.client.chat.completions.create(
                # model = "ep-20250207172624-lgfzb",  # your model endpoint ID
                model="ep-20250210135803-8r8cd",  # lz r1-满血
                # model="ep-20250210135919-kncgs",  # lz r1-32b
                messages=[
                    {"role": "user", "content": query_message},
                ],
                stream=True,
                timeout=120,
            )
            reasoning_content = ""
            content = ""
            async for chunk in stream:
                if hasattr(chunk.choices[0].delta, 'reasoning_content'):
                    print(chunk.choices[0].delta.reasoning_content, end='')
                    reasoning_content += chunk.choices[0].delta.reasoning_content
                    # yield json.dumps({'type': 'thinking', 'content': chunk.choices[0].delta.reasoning_content},
                    #                  ensure_ascii=False)
                else:
                    print(chunk.choices[0].delta.content, end='\n')
                    content += chunk.choices[0].delta.content
                    # yield json.dumps({'type': 'text', 'content': chunk.choices[0].delta.content},
                    #                  ensure_ascii=False)
            type_list = eval(Get_List(content))
            problem_type_index_list = sorted(list(set(type_list)))
            print('problem_type_index_list:', problem_type_index_list)

            if len(problem_type_index_list) == 0:
                print('problem_type_list:', problem_type_index_list)
                print('该招股书切片不适合被询问')
                problem_type_index_list = []
            else:
                if max(problem_type_index_list) <= 58 and min(problem_type_index_list) >= 0:
                    print('problem_type_list:', problem_type_index_list)
                else:
                    problem_type_index_list = []
                    raise ValueError('问题分类索引不在[0,58]范围内')
        except Exception as e:
            problem_type_index_list = []
            print(f"C5 error : {e}")
            # raise ValueError(f"C5 error : {e}")
            # yield f"C5 error : {e}"

        problem_type_name_list = []
        for type_index in problem_type_index_list:
            problem_type_name_list.append(problem_type_dict[type_index])
        print(problem_type_name_list)
        # problem_type_name_list = str(problem_type_name_list)

        e_time = time.time()
        print('单条时间', e_time - s_time)

        return problem_type_name_list


if __name__ == '__main__':

    import asyncio
    import threading

    bubble = Bubble()
    content_text = '''报告期各期末，公司应收账款账面价值分别为 11,999.50 万元、11128.76 万
元和 14.781.75 万元，'''
    sector = '科创板'
    # problem_type_name_list = asyncio.run(bubble.title_bubble(content_text, sector))
    problem_type_name_list = asyncio.run(bubble.title_bubble_v3(content_text, sector))
    # problem_type_name_list = asyncio.run(bubble.title_bubble_r1(content_text, sector))
    print(json.dumps(problem_type_name_list, ensure_ascii=False))

